package com.lry.pda.bottom.mysql.innodb;

import java.nio.charset.StandardCharsets;
import java.util.*;

public class BPlusTree<K extends Comparable<K>, V> {

    private int maxDegree;//每个节点的最大出度，如果为3，则节点内数据最多为2
    private TreeNode root;

    //默认为唯一主键索引
    private boolean uniqueKey = true;

    public BPlusTree(int maxDegree){
        this.maxDegree = maxDegree;
    }

    public BPlusTree(int maxDegree,boolean uniqueKey){
        this.maxDegree = maxDegree;
        this.uniqueKey = uniqueKey;
    }

    public void insert(K key,V value){
        if(null==root){
            root = new TreeLeafNode(asList(key),asList(asList(value)));
            return;
        }
        TreeNode treeNode = root.insert(key,value);
        if(null != treeNode){
            K rootKey = treeNode.getClass().equals(TreeLeafNode.class)
                    ? treeNode.keys.get(0)
                    : ((TreeDirectoryNode) treeNode).findLeafKey(treeNode);
            root = new TreeDirectoryNode(asList(rootKey),asList(root,treeNode));
        }
    }

    public List<V> query(K key){
        return root.query(key);
    }
    public List<V> rangeQuery(K startKey,K endKey){
        return root.rangeQuery(startKey,endKey);
    }
    public boolean update(K key,V value){
        return root.update(key,value);
    }

    private final <T> List<T> asList(T... e) {
        List<T> res = new ArrayList<>();

        Collections.addAll(res, e);
        return res;
    }

    @Override
    public String toString() {
        return root.toString();
    }

    private abstract class TreeNode {
        protected List<K> keys;
        public abstract TreeNode insert(K key,V value);
        public abstract List<V> query(K key);
        public abstract List<V> rangeQuery(K startKey,K endKey);
        public abstract boolean update(K key,V value);

        protected boolean isFull(){
            return keys.size() == maxDegree-1;
        }

        protected int getInsertIndex(K key) {
            int l = 0;
            int r = keys.size() - 1;
            int index = keys.size();
            while (l <= r) {
                int mid = l + ((r - l) >> 1);
                if (keys.get(mid).compareTo(key) > 0) {
                    index = mid;
                    r = mid - 1;
                } else {
                    l = mid + 1;
                }
            }
            return index;
        }

        protected int getQueryIndex(K key) {
            int l = 0;
            int r = keys.size() - 1;
            int index = keys.size();
            while (l <= r) {
                int mid = l + ((r - l) >> 1);
                if (keys.get(mid).compareTo(key) >= 0) {
                    index = mid;
                    r = mid - 1;
                } else {
                    l = mid + 1;
                }
            }
            return index;
        }
    }

    /**
     * 目录节点
     */
    private class TreeDirectoryNode extends TreeNode{
        private List<TreeNode> children;

        public TreeDirectoryNode(List<K> keys,List<TreeNode> children){
            this.keys = keys;
            this.children = children;
        }

        @Override
        public TreeNode insert(K key, V value) {
            TreeNode newNode = this.children.get(getInsertIndex(key)).insert(key, value);
            if(null==newNode){
                return null;
            }

            if(isFull()){
                TreeDirectoryNode newNonLeafNode = split(newNode);
                newNonLeafNode.keys.remove(0);
                return newNonLeafNode;
            }else{
                K leafKey = findLeafKey(newNode);
                int insertIndex = getInsertIndex(leafKey);
                this.keys.add(insertIndex,leafKey);
                this.children.add(insertIndex+1,newNode);
                return null;
            }
        }

        @Override
        public List<V> query(K key) {
            return this.children.get(getInsertIndex(key)).query(key);
        }

        @Override
        public List<V> rangeQuery(K startKey, K endKey) {
            return this.children.get(getInsertIndex(startKey)).rangeQuery(startKey,endKey);
        }

        @Override
        public boolean update(K key, V value) {
            return this.children.get(getInsertIndex(key)).update(key,value);
        }


        private TreeDirectoryNode split(TreeNode newNode) {
            K key = findLeafKey(newNode);
            List<K> allKeys = new ArrayList<>(this.keys);
            List<TreeNode> allChildren = new ArrayList<>(this.children);
            int insertIndex = getInsertIndex(key);
            allKeys.add(insertIndex,key);
            allChildren.add(insertIndex+1,newNode);
            int midIdx = allKeys.size()/2;

            List<K> rightKeys = new ArrayList<>(allKeys.subList(midIdx, allKeys.size()));
            List<TreeNode> rightData = new ArrayList<>(allChildren.subList(midIdx+1, allChildren.size()));
            TreeDirectoryNode newNonLeafNode = new TreeDirectoryNode(rightKeys, rightData);

            this.keys = new ArrayList<>(allKeys.subList(0, midIdx));
            this.children = new ArrayList<>(allChildren.subList(0, midIdx+1));

            return newNonLeafNode;
        }

        public K findLeafKey(TreeNode cur) {
            if (cur.getClass().equals(TreeLeafNode.class)) {
                return cur.keys.get(0);
            }
            return findLeafKey(((TreeDirectoryNode) cur).children.get(0));
        }

        @Override
        public String toString() {
            StringBuilder res = new StringBuilder();
            Queue<TreeNode> queue = new LinkedList<>();
            queue.add(this);
            while (!queue.isEmpty()) {
                int size = queue.size();
                for (int i = 0; i < size; ++i) {
                    TreeNode cur = queue.poll();
                    assert cur != null;
                    res.append(cur.keys).append("  ");
                    if (cur.getClass().equals(TreeDirectoryNode.class)) {
                        queue.addAll(((TreeDirectoryNode) cur).children);
                    }
                }
                res.append('\n');
            }
            return res.toString();
        }
    }

    /**
     * 叶子节点
     */
    private class TreeLeafNode extends TreeNode{
        private List<List<V>> data;
        private TreeLeafNode next;

        public TreeLeafNode(List<K> keys,List<List<V>> data){
            this.keys = keys;
            this.data = data;
        }

        @Override
        public TreeNode insert(K key, V value) {
            int index = getIndexByKey(key);
            if(index!=-1){
                if(uniqueKey){
                    throw new RuntimeException("唯一主键冲突");
                }else{
                    this.data.get(index).add(value);
                    return null;
                }
            }

            if(isFull()){
                //当前叶子节点放满了，需要进行页分裂
                TreeLeafNode newLeafNode = split(key,value);
                return newLeafNode;
            }else{
                index = getInsertIndex(key);
                keys.add(index,key);
                data.add(index,asList(value));
                return null;
            }
        }

        @Override
        public List<V> query(K key) {
            int idx = getIndexByKey(key);
            if(!dataIdxCheck(idx)){
                return Collections.emptyList();
            }
            return this.data.get(idx);
        }

        @Override
        public List<V> rangeQuery(K startKey, K endKey) {
            int startIdx = getQueryIndex(startKey);
            int endIdx = getQueryIndex(endKey);

            List<V> ansList = new ArrayList<>();
            for(int i=startIdx;i<endIdx;i++){
                ansList.addAll(this.data.get(i));
            }

            TreeLeafNode nextNode = this.next;
            while(null!=nextNode){
                for (int i = 0; i < nextNode.keys.size(); ++i) {
                    if (nextNode.keys.get(i).compareTo(endKey) <= 0) {
                        ansList.addAll(nextNode.data.get(i));
                    } else {
                        return ansList;
                    }
                }
                nextNode = nextNode.next;
            }
            return ansList;
        }

        @Override
        public boolean update(K key, V value) {
            int idx = getIndexByKey(key);
            if(!dataIdxCheck(idx)){
                return false;
            }
            this.data.get(idx).clear();
            this.data.get(idx).add(value);
            return true;
        }

        private boolean dataIdxCheck(int idx){
            if(idx<0 || idx>=data.size()){
                return false;
            }
            return true;
        }

        private TreeLeafNode split(K key,V value) {
            List<K> allKeys = new ArrayList<>(this.keys);
            List<List<V>> allData = new ArrayList<>(this.data);
            int insertIndex = getInsertIndex(key);
            allKeys.add(insertIndex,key);
            allData.add(insertIndex,asList(value));
            int midIdx = allKeys.size()/2;

            List<K> rightKeys = new ArrayList<>(allKeys.subList(midIdx, allKeys.size()));
            List<List<V>> rightData = new ArrayList<>(allData.subList(midIdx, allData.size()));
            TreeLeafNode newLeafNode = new TreeLeafNode(rightKeys, rightData);

            this.keys = new ArrayList<>(allKeys.subList(0, midIdx));
            this.data = new ArrayList<>(allData.subList(0, midIdx));

            newLeafNode.next = this.next;
            this.next = newLeafNode;
            return newLeafNode;
        }

        private int getIndexByKey(K key) {
            int l = 0;
            int r = keys.size() - 1;
            while (l <= r) {
                int mid = l + ((r - l) >> 1);
                int compare = keys.get(mid).compareTo(key);
                if (compare == 0) {
                    return mid;
                } else if (compare > 0) {
                    r = mid - 1;
                } else {
                    l = mid + 1;
                }
            }
            return -1;
        }

        @Override
        public String toString() {
            return keys.toString();
        }
    }


}
